import os
import time
import asyncio
from langchain_groq import ChatGroq
from langchain.schema import HumanMessage, SystemMessage
from config.settings import GROQ_API_KEY

# 初始化Groq客户端
llm = ChatGroq(
    temperature=0.3,
    groq_api_key=GROQ_API_KEY,
    model_name="llama-3.3-70b-versatile"  # 使用当前支持的模型
)

def summarize_text(content: str, language: str = "Chinese", max_retries: int = 3, base_delay: int = 60, max_wait_time: int = 30) -> str:
    """
    使用Groq API生成指定语言的摘要，包含重试机制
    
    Args:
        content: 要总结的文本内容
        language: 目标语言 ("Chinese" 或 "English")
        max_retries: 最大重试次数
        base_delay: 基础延迟时间（秒）
        max_wait_time: 最大等待时间（秒），超过此时间将跳过重试
    
    Returns:
        生成的摘要文本
    """
    if language == "Chinese":
        # 中文提示词：面向研究人员的技术性摘要生成
        # 目标：生成高质量的学术摘要，突出技术创新和方法论贡献
        system_prompt = """你是一位资深的AI/ML研究领域专家和学术文献分析师。请为给定的学术论文生成一个面向研究人员的技术性中文摘要。

技术要求：
1. 摘要长度：200-300字，确保信息密度和技术深度
2. 核心内容：
   - 明确指出技术创新点和方法论贡献
   - 详细描述所提出的算法、模型或架构
   - 量化实验结果和性能指标（如有）
   - 与现有方法的技术对比和优势
3. 技术表达：
   - 使用准确的学术术语和技术概念
   - 保留关键的英文术语（如模型名称、算法名称）
   - 突出技术细节和实现方法
   - 面向具有相关背景的研究人员
4. 结构化表述：
   - 问题定义与动机
   - 技术方法与创新
   - 实验验证与结果
   - 学术价值与影响

请生成专业的技术摘要，无需额外格式或标题。"""
        
        user_prompt = f"请为以下学术论文生成面向研究人员的技术性中文摘要：\n\n{content[:3000]}"
        
    else:  # English
        # 英文提示词：面向研究人员的技术性摘要生成
        # 目标：生成符合学术标准的英文摘要，强调技术贡献和方法创新
        system_prompt = """You are a senior AI/ML research expert and academic literature analyst. Please generate a technical English summary for the given academic paper, targeted at researchers and practitioners in the field.

Technical Requirements:
1. Summary length: 150-200 words with high information density and technical depth
2. Core content focus:
   - Clearly articulate technical innovations and methodological contributions
   - Describe proposed algorithms, models, or architectures in detail
   - Include quantitative experimental results and performance metrics (when available)
   - Compare with existing methods and highlight technical advantages
3. Technical expression:
   - Use precise academic terminology and technical concepts
   - Maintain rigorous scientific language appropriate for peer review
   - Emphasize technical details and implementation approaches
   - Target audience: researchers with relevant domain expertise
4. Structured presentation:
   - Problem formulation and motivation
   - Technical approach and innovations
   - Experimental validation and results
   - Academic significance and impact

Generate a professional technical summary without additional formatting or titles."""
        
        user_prompt = f"Please generate a technical English summary for researchers based on the following academic paper:\n\n{content[:3000]}"
    
    for attempt in range(max_retries):
        try:
            messages = [
                SystemMessage(content=system_prompt),
                HumanMessage(content=user_prompt)
            ]
            
            response = llm.invoke(messages)
            return response.content.strip()
            
        except Exception as e:
            error_str = str(e)
            print(f"[ERROR] Attempt {attempt + 1}/{max_retries} failed: {error_str}")
            
            # 检查是否是速率限制错误
            if "rate_limit_exceeded" in error_str or "429" in error_str:
                if attempt < max_retries - 1:  # 不是最后一次尝试
                    # 从错误消息中提取等待时间，如果无法提取则使用默认延迟
                    wait_time = base_delay
                    if "Please try again in" in error_str:
                        try:
                            # 尝试从错误消息中提取等待时间
                            import re
                            time_match = re.search(r'Please try again in (\d+)m(\d+(?:\.\d+)?)s', error_str)
                            if time_match:
                                minutes = int(time_match.group(1))
                                seconds = float(time_match.group(2))
                                wait_time = minutes * 60 + seconds + 10  # 额外加10秒缓冲
                        except:
                            pass
                    
                    # 检查等待时间是否超过最大限制
                    if wait_time > max_wait_time:
                        error_msg = f"Rate limit wait time ({wait_time}s) exceeds maximum allowed time ({max_wait_time}s). Skipping {language} summary generation."
                        print(f"[WARNING] {error_msg}")
                        return error_msg
                    
                    print(f"[INFO] Rate limit hit. Waiting {wait_time} seconds before retry...")
                    time.sleep(wait_time)
                    continue
            else:
                # 非速率限制错误，等待较短时间后重试
                if attempt < max_retries - 1:
                    wait_time = (attempt + 1) * 5  # 递增延迟：5s, 10s, 15s...
                    print(f"[INFO] Waiting {wait_time} seconds before retry...")
                    time.sleep(wait_time)
                    continue
        
        # 如果是最后一次尝试或者其他错误，返回错误信息
        if attempt == max_retries - 1:
            error_msg = f"Failed to generate {language} summary after {max_retries} attempts: {str(e)}"
            print(f"[ERROR] {error_msg}")
            return error_msg

async def main():
    """
    主入口函数，用于测试双语摘要功能。
    """
    if not settings.GROQ_API_KEY:
        print("请在.env文件中或通过环境变量设置 GROQ_API_KEY 以运行此测试。")
        return

    sample_text = """
    Artificial intelligence (AI) is rapidly changing the world. From self-driving cars to medical diagnostics, the applications of AI are widespread. 
    Large language models (LLMs) are a significant breakthrough in the field of AI, capable of understanding and generating human-like text. 
    This has provided powerful capabilities for natural language processing tasks such as summarization, translation, and question-answering. 
    However, the development of AI also brings challenges, including data privacy, algorithmic bias, and changes in the job market. 
    We need to address these societal issues while promoting technological advancement.
    """ * 3

    print("正在使用Groq生成双语摘要...")
    summaries = await summarize_text(sample_text)
    
    print("\n--- English Summary ---")
    print(summaries["summary_en"])
    
    print("\n--- 中文摘要 ---")
    print(summaries["summary_zh"])

if __name__ == '__main__':
    asyncio.run(main())